Add support for float8 activation for Int4PreshuffledTensor #2437

jerryzh168 · 2025-06-24T22:25:42Z

Stacked PRs:

Add support for float8 activation for Int4PreshuffledTensor

Summary:
Added basic op support like linear and bmm, we have both float8 and bf16 in the same Tensor
because it's the same dtype, only difference is whether the activation is quantized or not. Although
there is some differneces in implementation:

bf16 activaton:

group_scale
group_zero

fp8 activation

group_scale
row_scale

Test Plan:
python test/quantization/quantize_/workflows/int4/test_int4_preshuffled_tensor.py

Reviewers:

Subscribers:

Tasks:

Tags:

pytorch-bot · 2025-06-24T22:25:46Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/2437

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (1 Unrelated Failure)

As of commit 889dca3 with merge base e5ca515 ():

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh) (trunk failure)
##[error]The operation was canceled.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

Summary: Note: slice is not working yet, others are working Test Plan: python test/dtypes/test_float8_activation_int4_groupwise_preshuffle.py Reviewers: Subscribers: Tasks: Tags: stack-info: PR: #2437, branch: jerryzh168/stack/4

test/quantization/quantize_/int4/test_int4_groupwise_preshuffle_tensor.py

torchao/quantization/quant_api.py

torchao/quantization/quantize_/int4/int4_groupwise_preshuffle_tensor.py

drisspg

Can you add serialization tests

jerryzh168 · 2025-07-08T02:02:07Z

Can you add serialization tests

you mean serialization of the models? it's moved to https://github.com/pytorch/ao/pull/2463/files#diff-9f6b6c4b39656e797cfda97536a4cf8a82004c64da518ad524637b471b716739, I don't exactly remember the reason

for config serialization we did a config refactor in the last PR, I can add after we are aligned on what config should look like

torchao/quantization/quantize_/int4/__init__.py

torchao/quantization/quant_api.py

torchao/quantization/quantize_/int4/int4_groupwise_preshuffle_tensor.py

Summary: Added basic op support like linear and bmm, we have both float8 and bf16 in the same Tensor because it's the same dtype, only difference is whether the activation is quantized or not. Although there is some differneces in implementation: bf16 activaton: * group_scale * group_zero fp8 activation * group_scale * row_scale Test Plan: python test/quantization/quantize_/workflows/int4/test_int4_preshuffled_tensor.py Reviewers: Subscribers: Tasks: Tags: stack-info: PR: #2437, branch: jerryzh168/stack/4

jerryzh168 force-pushed the jerryzh168/stack/4 branch from d6d3477 to 26517e8 Compare June 24, 2025 22:25

This was referenced Jun 24, 2025

Add support for Int4GroupwisePreshuffleTensor for fbgemm #2421

Merged

Remove transpose_input from fbgemm configs #2422

Merged

facebook-github-bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Jun 24, 2025

jerryzh168 changed the base branch from jerryzh168/stack/2 to main June 24, 2025 22:26

jerryzh168 force-pushed the jerryzh168/stack/4 branch from 26517e8 to d187f78 Compare June 24, 2025 22:26

jerryzh168 changed the base branch from main to jerryzh168/stack/2 June 24, 2025 22:26

jerryzh168 changed the base branch from jerryzh168/stack/2 to main June 24, 2025 22:28

jerryzh168 force-pushed the jerryzh168/stack/4 branch from d187f78 to 2fcff42 Compare June 24, 2025 22:28

jerryzh168 changed the base branch from main to jerryzh168/stack/2 June 24, 2025 22:28

jerryzh168 changed the base branch from jerryzh168/stack/2 to main June 26, 2025 05:03

jerryzh168 force-pushed the jerryzh168/stack/4 branch from 2fcff42 to 95856ed Compare June 26, 2025 05:03

jerryzh168 changed the base branch from main to jerryzh168/stack/2 June 26, 2025 05:03

jerryzh168 changed the base branch from jerryzh168/stack/2 to main June 27, 2025 19:36

jerryzh168 force-pushed the jerryzh168/stack/4 branch from 95856ed to 1dec2cb Compare June 27, 2025 19:37

jerryzh168 changed the base branch from main to jerryzh168/stack/2 June 27, 2025 19:37

jerryzh168 changed the base branch from jerryzh168/stack/2 to main June 27, 2025 19:38

jerryzh168 force-pushed the jerryzh168/stack/4 branch from 1dec2cb to 1645c79 Compare June 27, 2025 19:38

jerryzh168 changed the base branch from main to jerryzh168/stack/2 June 27, 2025 19:38

jerryzh168 changed the base branch from jerryzh168/stack/2 to main June 27, 2025 19:48

jerryzh168 force-pushed the jerryzh168/stack/4 branch from 1645c79 to 5e9e869 Compare June 27, 2025 19:48

jerryzh168 changed the base branch from main to jerryzh168/stack/2 June 27, 2025 19:48

jerryzh168 changed the base branch from main to jerryzh168/stack/2 July 3, 2025 02:38

jerryzh168 changed the base branch from jerryzh168/stack/2 to main July 3, 2025 02:44

jerryzh168 changed the base branch from main to jerryzh168/stack/2 July 3, 2025 02:44

jerryzh168 changed the base branch from jerryzh168/stack/2 to main July 3, 2025 20:58

jerryzh168 force-pushed the jerryzh168/stack/4 branch 5 times, most recently from 72bc113 to ea16397 Compare July 7, 2025 19:57